// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/functions/prestosql/array_constructor.h>
#include <pollux/expression/expr.h>
#include <pollux/expression/vector_function.h>

namespace kumo::pollux::functions {
    namespace {
        class ArrayConstructor : public exec::VectorFunction {
        public:
            void apply(
                const SelectivityVector &rows,
                std::vector<VectorPtr> &args,
                const TypePtr &outputType,
                exec::EvalCtx &context,
                VectorPtr &result) const override {
                auto numArgs = args.size();

                context.ensure_writable(rows, outputType, result);
                result->clear_nulls(rows);
                auto arrayResult = result->as<ArrayVector>();
                auto sizes = arrayResult->mutableSizes(rows.end());
                auto rawSizes = sizes->asMutable<int32_t>();
                auto offsets = arrayResult->mutableOffsets(rows.end());
                auto rawOffsets = offsets->asMutable<int32_t>();

                auto elementsResult = arrayResult->elements();

                // append to the end of the "elements" vector
                auto baseOffset = elementsResult->size();

                if (args.empty()) {
                    rows.applyToSelected([&](vector_size_t row) {
                        rawSizes[row] = 0;
                        rawOffsets[row] = baseOffset;
                    });
                } else {
                    elementsResult->resize(baseOffset + numArgs * rows.countSelected());

                    if (shouldCopyRanges(elementsResult->type())) {
                        std::vector<BaseVector::CopyRange> ranges;
                        ranges.reserve(rows.end());

                        vector_size_t offset = baseOffset;
                        rows.applyToSelected([&](vector_size_t row) {
                            rawSizes[row] = numArgs;
                            rawOffsets[row] = offset;
                            ranges.push_back({row, offset, 1});
                            offset += numArgs;
                        });

                        elementsResult->copy_ranges(args[0].get(), ranges);

                        for (int i = 1; i < numArgs; i++) {
                            for (auto &range: ranges) {
                                ++range.targetIndex;
                            }
                            elementsResult->copy_ranges(args[i].get(), ranges);
                        }
                    } else {
                        SelectivityVector targetRows(elementsResult->size(), false);
                        std::vector<vector_size_t> toSourceRow(elementsResult->size());

                        vector_size_t offset = baseOffset;
                        rows.applyToSelected([&](vector_size_t row) {
                            rawSizes[row] = numArgs;
                            rawOffsets[row] = offset;

                            targetRows.setValid(offset, true);
                            toSourceRow[offset] = row;

                            offset += numArgs;
                        });
                        targetRows.updateBounds();
                        elementsResult->copy(args[0].get(), targetRows, toSourceRow.data());

                        for (int i = 1; i < numArgs; i++) {
                            targetRows.clearAll();

                            vector_size_t offset_2 = baseOffset;
                            rows.applyToSelected([&](vector_size_t row) {
                                targetRows.setValid(offset_2 + i, true);
                                toSourceRow[offset_2 + i] = row;
                                offset_2 += numArgs;
                            });

                            targetRows.updateBounds();
                            elementsResult->copy(args[i].get(), targetRows, toSourceRow.data());
                        }
                    }
                }
            }

            static std::vector<std::shared_ptr<exec::FunctionSignature> > signatures() {
                return {
                    // () -> array(unknown)
                    exec::FunctionSignatureBuilder().returnType("array(unknown)").build(),
                    // T... -> array(T)
                    exec::FunctionSignatureBuilder()
                    .typeVariable("T")
                    .returnType("array(T)")
                    .variableArity("T")
                    .build(),
                };
            }

        private:
            // BaseVector::copyRange is faster for arrays and maps and slower for
            // primitive types. Check if 'type' is an array or map or contains an array or
            // map. If so, return true, otherwise, false.
            static bool shouldCopyRanges(const TypePtr &type) {
                if (type->isPrimitiveType()) {
                    return false;
                }

                if (!type->is_row()) {
                    return true;
                }

                const auto &rowType = type->as_row();
                for (const auto &child: rowType.children()) {
                    if (shouldCopyRanges(child)) {
                        return true;
                    }
                }
                return false;
            }
        };
    } // namespace

    POLLUX_DECLARE_VECTOR_FUNCTION_WITH_METADATA(
        udf_array_constructor,
        ArrayConstructor::signatures(),
        exec::VectorFunctionMetadataBuilder().defaultNullBehavior(false).build(),
        std::make_unique<ArrayConstructor>());

    void registerArrayConstructor(const std::string &name) {
        POLLUX_REGISTER_VECTOR_FUNCTION(udf_array_constructor, name);
    }
} // namespace kumo::pollux::functions
